"""
run on client-backend.
read insert tasks from redis, process, update task state lists.
read grib file received by client-backend and insert into influxdb on respective datanode.
"""

import datetime
import logging
import os
import time

import pygrib
import redis
import requests
import zstd

import logger as my_log
import series_key

# redisKeyNewTaskId incremental id for next uploaded grib file(and task id)
redisKeyNewTaskId = "client-backend:insert:new-task-id"

# redisKeyTmpFileSet set of name of tmp grib files
# redisKeyTmpFileSet = "client-backend:insert:tmp-file"

# redisKeyInsertTaskListAll [set]task id list
redisKeyInsertTaskListAll = "client-backend:insert:task-list:all"

# redisKeyInsertTaskListDispatched hash(task_id:heartbeat) of dispatched tasks, need to be watched
redisKeyInsertTaskListDispatched = "client-backend:insert:task-list:dispatched"

# redisKeyInsertTaskListWait list(task_id) of task need to dispatch
redisKeyInsertTaskListWait = "client-backend:insert:task-list:wait"

# redisKeyInsertTaskStateMessage hash(task_id:json_str_message) for task state
# keys including size, progress, status message
redisKeyInsertTaskStateMessage = "client-backend:insert:task-state:message"

my_log.init()
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

rdb = redis.Redis()
redisGetTaskStr = """
    -- KEYS[1]: task-list:all
    -- KEYS[2]: task-list:dispatched
    -- KEYS[3]: task-state:message
    -- ARGV[1]: task-id

    local stateStr = redis.call("hget", KEYS[3], ARGV[1])
    if stateStr == false
    then
        -- return when task is deleted
        return ""
    end
    local stateObject = cjson.decode(stateStr)
    local state = stateObject["state"]
    if state ~= 0
    then
        -- task state changed by other routine(may be paused by user)
        -- insert executor need to fetch other tasks
        return ""
    else
        -- when state is set to running by this routine, task can be executed.
        stateObject["state"] = 1
        stateObject["message"] = "task has been dispatched"
        local tmp_grib_path = stateObject["tmp_grib_file"]
        
        stateStr = cjson.encode(stateObject)
        redis.call("hset", KEYS[3], ARGV[1], stateStr)
        
        return tmp_grib_path
    end
"""
redisGetTask = rdb.register_script(redisGetTaskStr)

# ignore pause
redisUpdateStateStr = """
    -- KEYS[1]: task-list:all
    -- KEYS[2]: task-list:dispatched
    -- KEYS[3]: task-state:message
    -- ARGV[1]: task-id
    -- ARGV[2]: progress
    -- ARGV[3]: message
    -- ARGV[4]: state
    
    if redis.call("SISMEMBER", KEYS[1], ARGV[1]) == 1
    then
        local stateStr = redis.call("hget", KEYS[3], ARGV[1])
        local stateObject = cjson.decode(stateStr)
        stateObject["progress"] = tonumber(ARGV[2])
        stateObject["message"] = ARGV[3]
        stateObject["state"] = tonumber(ARGV[4])
        stateStr = cjson.encode(stateObject)
        redis.call("hset", KEYS[3], ARGV[1], stateStr)
        
        
        
        return "continue"
    else
        return "delete"
    end

"""
redisUpdateState = rdb.register_script(redisUpdateStateStr)


def insert(grib_file_path: str, task_id: str) -> None:
    if not os.path.isfile(grib_file_path):
        logger.error(f"file not exists {grib_file_path}")
        return

    file_start_time = time.time()

    msg_len = 0

    # get series_key of all messages
    series_key_list = []
    meta_info_list = []
    with pygrib.open(grib_file_path) as grib_file:
        msg_len = grib_file.messages
        for msg in grib_file:
            msg_start_time = time.time()

            values_count = msg["getNumberOfValues"]
            if values_count != 1038240:
                logger.error(f"incomplete data")
            measurement = msg["parameterName"]
            pressure_level = str(msg["level"])
            timestamp = datetime.datetime(msg["year"], msg["month"], msg["day"],
                                          msg["hour"], msg["minute"], msg["second"],
                                          tzinfo=datetime.timezone.utc).timestamp()
            timestamp = int(timestamp)
            meta_info_list.append({
                'variable': measurement,
                'pressure_level': pressure_level,
                'timestamp': timestamp,
            })
            series_key_list.append(series_key.get_series_key(measurement, pressure_level, timestamp))

    generated_key_time = time.time()
    logger.info(f"generate key used {generated_key_time - file_start_time:.2f}")

    # send series_key to name_node, query data_node IP
    unique_key_list = list(set(series_key_list))
    upload_keys = bytes()
    for key in unique_key_list:
        upload_keys += key
    # url = "http://10.249.177.55:8081/locate"
    url = "http://localhost:8081/locate"
    response = requests.post(url, files={'data': upload_keys})
    if response.status_code != 200:
        err = f"failed to get datanode from namenode: [{response.status_code}] {response.text}"
        logger.error(err)
        raise Exception(err)
    data_nodes_list = response.text.strip().split(",")
    key_node_map = {}
    for key, node_ip in zip(unique_key_list, data_nodes_list):
        key_node_map[key] = node_ip
    got_data_node_time = time.time()

    action = ""
    # upload messages to different data node
    with pygrib.open(grib_file_path) as grib_file:
        for i, msg in enumerate(grib_file):
            progress = int(i / msg_len * 100)
            state = 1
            data_node_ip = key_node_map[series_key_list[i]]
            block_info = meta_info_list[i]
            message = f"insert {block_info} into {data_node_ip}"

            action = redisUpdateState(
                keys=[redisKeyInsertTaskListAll,
                      redisKeyInsertTaskListDispatched,
                      redisKeyInsertTaskStateMessage],
                args=[task_id, progress, message, state]).decode()
            logger.info(action)
            if action == "delete":
                logger.info(f"task deleted before finish")
                break
            logger.info(message)
            upload_start_time = time.time()

            compressed_data = zstd.compress(msg.values, 1)
            upload_params = block_info.copy()
            upload_params["variable"] = series_key.variable_name2id[upload_params["variable"]]
            upload_params["pressure_level"] = series_key.pressure_name2id[upload_params["pressure_level"]]
            resp = requests.post(f"http://{data_node_ip}:8082/insert", params=upload_params,
                                 files={"data": compressed_data})
            if resp.status_code != 200:
                err = f"failed to insert {block_info}"
                logger.error(err)
                raise Exception(err)
            else:
                logger.info(f"insert {block_info} success")
            logger.info(f"insert used {time.time() - upload_start_time:.2f} seconds")
    finish_task_time = time.time()

    if action == "delete":
        os.remove(grib_file_path)
    elif action == "continue":
        os.remove(grib_file_path)
        message = f"insert finished, used {finish_task_time - got_data_node_time:.2f} seconds"
        logger.info(message)
        state = 3
        progress = 100
        redisUpdateState(
            keys=[redisKeyInsertTaskListAll,
                  redisKeyInsertTaskListDispatched,
                  redisKeyInsertTaskStateMessage],
            args=[task_id, progress, message, state])


def main():
    while True:
        task_id = rdb.blpop(redisKeyInsertTaskListWait)
        task_id = task_id[1].decode()
        logger.info(f"got task {task_id}")
        grib_file_path = redisGetTask(
            keys=[redisKeyInsertTaskListAll, redisKeyInsertTaskListDispatched,
                  redisKeyInsertTaskStateMessage],
            args=[task_id]).decode()
        if grib_file_path == "":
            continue
        logger.info(f"got task {task_id}")
        try:
            insert(grib_file_path, task_id)
        except Exception as err:
            message = f"insert failed, err: {err}"
            logger.error(message)
            state = 2
            progress = 0
            redisUpdateState(
                keys=[redisKeyInsertTaskListAll,
                      redisKeyInsertTaskListDispatched,
                      redisKeyInsertTaskStateMessage],
                args=[task_id, progress, message, state])


if __name__ == '__main__':
    main()
